import pickle as pkl
from query import Query
import numpy as np
import joblib
import os
class sample:
    def __init__(self, active, fn, club_member_status,\
        fashion_news_frequency, age, pcode, sales_channel_id, 
        art_seq, click):
        self.sparse_feature = [active,fn,club_member_status, 
        fashion_news_frequency,pcode,sales_channel_id]
        self.dense_feature = [age]
        self.tag = click
        self.historyfeature = []
        for i in art_seq:
            self.historyfeature.append(i)
    def sample2np(self):
        ret = []
        return ret


class batch_getter:
    def __init__(self, i, use_neg = True, neg_sample = 100):
        self.Query = Query(i)
        self.use_neg = use_neg
        self.neg_sample = neg_sample
        word_emb_tenor_file = open('./Data/Processed/word_emb_tensor.pkl', 'rb')
        topk_index_100_file = open('./Data/Processed/topk_index_100.pkl','rb')
        article_sparse_file = open('./Data/Processed/article_sparse.pkl','rb')
        self.artid = pkl.load(article_sparse_file)
        self.pcode = pkl.load(article_sparse_file)
        self.word_emb_tensor = pkl.load(word_emb_tenor_file)
        self.topk_index_100 = pkl.load(topk_index_100_file)
        print('loaded!')
    
    def get_batch(self, N = 100):
        Query_feature = self.Query.query(N)
        N = Query_feature['length']
        #word_feature = np.array([self.word_emb_tensor[i - 1] for i in Query_feature['pos_artid']])
        #Query_feature["word_feature"] = word_feature
        # negQueryfeature = {}
        # for i in Query_feature.keys():
        #     negQueryfeature[i] = np.array([])
        # negQueryfeature['length'] = 0
        # for x, y in Query_feature.items():
        #     print(x, y.shape)
        if self.use_neg:
            for i in range(N):
                neg_pos_custom_id = np.array([Query_feature['pos_custom_id'][i] for j in range(self.neg_sample)])
                neg_pos_FN = np.array([Query_feature['pos_FN'][i] for j in range(self.neg_sample)])
                neg_pos_Active = np.array([Query_feature['pos_Active'][i] for j in range(self.neg_sample)])
                neg_pos_club_member_status = np.array([Query_feature['pos_club_member_status'][i] for j in range(self.neg_sample)])
                neg_pos_fashion_news_frequency = np.array([Query_feature['pos_fashion_news_frequency'][i] for j in range(self.neg_sample)])
                neg_pos_age = np.array([Query_feature['pos_age'][i] for j in range(self.neg_sample)])
                neg_pos_artid = np.array([self.topk_index_100[Query_feature['pos_artid'][i] - 1][j] + 1 for j in range(self.neg_sample)])
                neg_pos_pcode = np.array([self.pcode[self.topk_index_100[Query_feature['pos_artid'][i] - 1][j]] for j in range(self.neg_sample)])
                #neg_word_feature = np.array([self.word_emb_tensor[self.topk_index_100[Query_feature['pos_artid'][i]][j]] for j in range(self.neg_sample)])
                neg_pos_sales_channel_id = np.array([Query_feature['pos_sales_channel_id'][i] for j in range(self.neg_sample)])
                neg_pos_pcode_seq = np.array([Query_feature['hist_pos_pcode'][i] for j in range(self.neg_sample)])
                neg_pos_artid_seq = np.array([Query_feature['hist_pos_artid'][i] for j in range(self.neg_sample)])
                neg_pos_click = np.array([0 for j in range(self.neg_sample)])
                neg_seq_length = np.array([Query_feature['seq_length'][i] for j in range(self.neg_sample)])
                Query_feature['pos_custom_id'] = np.concatenate((Query_feature['pos_custom_id'],neg_pos_custom_id))
                Query_feature['pos_FN'] = np.concatenate((Query_feature['pos_FN'],neg_pos_FN))
                Query_feature['pos_Active'] = np.concatenate((Query_feature['pos_Active'],neg_pos_Active))
                Query_feature['pos_club_member_status'] = np.concatenate((Query_feature['pos_club_member_status'],neg_pos_club_member_status))
                Query_feature['pos_fashion_news_frequency'] = np.concatenate((Query_feature['pos_fashion_news_frequency'],neg_pos_fashion_news_frequency))
                Query_feature['pos_age'] = np.concatenate((Query_feature['pos_age'],neg_pos_age))
                Query_feature['pos_artid'] = np.concatenate((Query_feature['pos_artid'],neg_pos_artid))
                Query_feature['pos_pcode'] = np.concatenate((Query_feature['pos_pcode'],neg_pos_pcode))
                #Query_feature['word_feature'] = np.concatenate((Query_feature['word_feature'],neg_word_feature))
                Query_feature['pos_sales_channel_id'] = np.concatenate((Query_feature['pos_sales_channel_id'],neg_pos_sales_channel_id))
                Query_feature['hist_pos_pcode'] = np.concatenate((Query_feature['hist_pos_pcode'],neg_pos_pcode_seq))
                Query_feature['hist_pos_artid'] = np.concatenate((Query_feature['hist_pos_artid'],neg_pos_artid_seq))
                Query_feature['seq_length'] = np.concatenate((Query_feature['seq_length'],neg_seq_length))
                Query_feature['pos_click'] = np.concatenate((Query_feature['pos_click'],neg_pos_click))
                Query_feature['length'] += self.neg_sample
                # negQueryfeature['pos_custom_id'] = np.concatenate((negQueryfeature['pos_custom_id'],neg_pos_custom_id))
                # negQueryfeature['pos_FN'] = np.concatenate((negQueryfeature['pos_FN'],neg_pos_FN))
                # negQueryfeature['pos_Active'] = np.concatenate((negQueryfeature['pos_Active'],neg_pos_Active))
                # negQueryfeature['pos_club_member_status'] = np.concatenate((negQueryfeature['pos_club_member_status'],neg_pos_club_member_status))
                # negQueryfeature['pos_fashion_news_frequency'] = np.concatenate((negQueryfeature['pos_fashion_news_frequency'],neg_pos_fashion_news_frequency))
                # negQueryfeature['pos_age'] = np.concatenate((negQueryfeature['pos_age'],neg_pos_age))
                # negQueryfeature['pos_artid'] = np.concatenate((negQueryfeature['pos_artid'],neg_pos_artid))
                # negQueryfeature['pos_pcode'] = np.concatenate((negQueryfeature['pos_pcode'],neg_pos_pcode))
                # negQueryfeature['word_feature'] = np.concatenate((negQueryfeature['word_feature'],neg_word_feature))
                # negQueryfeature['pos_sales_channel_id'] = np.concatenate((negQueryfeature['pos_sales_channel_id'],neg_pos_sales_channel_id))
                # negQueryfeature['hist_pos_pcode'] = np.concatenate((negQueryfeature['hist_pos_pcode'],neg_pos_pcode_seq))
                # negQueryfeature['hist_pos_artid'] = np.concatenate((negQueryfeature['hist_pos_artid'],neg_pos_artid_seq))
                # negQueryfeature['seq_length'] = np.concatenate((negQueryfeature['seq_length'],neg_seq_length))
                # negQueryfeature['pos_click'] = np.concatenate((negQueryfeature['pos_click'],neg_pos_click))
                # negQueryfeature['length'] += self.neg_sample
        return Query_feature#,negQueryfeature


        


        